# %%
# nn模块
# 使用nn模块定义4阶多项式来拟合sin(x)
import torch
import math

# 创建数据集
x = torch.linspace(-math.pi, math.pi, 20000)
y = torch.sin(x)

# pytorch 文档原文:
# For this example, the output y is a linear function of (x, x^2, x^3)
# 为什么这个函数是线性的呢？
# 不是非线性吗?仔细想一下,我们求梯度是loss对参数求,和x,y实际上是没什么关系的
# 因此用线性层神经网络，创建好tensor(x,x^2,x^3)
p = torch.tensor([1, 2, 3])
xx = x.unsqueeze(-1).pow(p)
# x的shape是(2000),unsqueeze(-1)后的shape是(2000,1) p的shape是(3,)
# 经过广播得到xx的shape是(2000,3)
# nn.Sequential是装nn的容器,
model = torch.nn.Sequential(
    torch.nn.Linear(3, 1),  # input_shape:(N,*,3) output_shape:(N,*,1)
    torch.nn.Flatten(0, 1)
)
# nn包也包含有多种损失函数,在这里我们用MSE
# reduction='sum'求了MSE之后用求和的方式变成1维
loss_fn = torch.nn.MSELoss(reduction='sum')

# 训练
learning_rate = 1e-7
for t in range(2000):
    #  Module objects override the __call__ operator
    # 可以直接像函数一样调用
    y_pred = model(xx)
    # 计算并打印loss,输入y_pred和y
    loss = loss_fn(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())
    # 梯度置零 然后反向传播
    model.zero_grad()
    # 反向传播对所有需要计算梯度的参数进行计算
    loss.backward()
    # 用梯度下降更新权重
    # 更新权重的计算不需要被保护到计算图里面
    with torch.no_grad():
        # 挨个从模型中拿出参数进行梯度下降更新
        for param in model.parameters():
            param -= learning_rate * param.grad
# 可以像访问列表一样访问模型的层
linear_layer = model[0]
print(linear_layer)
print(f'Result: y = {linear_layer.bias.item()} + '
      f'{linear_layer.weight[:, 0].item()} x + '
      f'{linear_layer.weight[:, 1].item()} x^2 +'
      f' {linear_layer.weight[:, 2].item()} x^3')
# %%
# 更进一步,不用torch.no_grad()更新,而是用optim更新
import torch
import math

# 相同的部分我就写注释了
x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)

p = torch.tensor([1, 2, 3])
xx = x.unsqueeze(-1).pow(p)

model = torch.nn.Sequential(
    torch.nn.Linear(3, 1),
    torch.nn.Flatten(0, 1)
)
loss_fn = torch.nn.MSELoss(reduction='sum')

# 用otim包来定义优化器更新,这里用RMSprop
learning_rate = 1e-3
# 记住一定要给优化器传入模型的参数
optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate)
for t in range(2000):
    y_pred = model(xx)

    loss = loss_fn(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # 这里用optimizer来清除梯度zero_grad
    optimizer.zero_grad()
    # 反向传播计算loss关于模型参数的梯度
    loss.backward()
    # 用step方法来更新梯度
    optimizer.step()

linear_layer = model[0]
print(f'Result: y = {linear_layer.bias.item()} + '
      f'{linear_layer.weight[:, 0].item()} x + '
      f'{linear_layer.weight[:, 1].item()} x^2 + '
      f'{linear_layer.weight[:, 2].item()} x^3')
# %%
# 自定义nn模块
# 如何更加自由的搭建nn模块？通过子类化nn.Module,然后自定义forward来实现
# 仍然是三阶多项式的例子
import torch
import math


class Polynomial3(torch.nn.Module):
    def __init__(self):
        """ 把需要用的参数或者模型层实例化 """
        # """"初始化超类""""
        super(Polynomial3, self).__init__()
        # 先把把需要用到的模型参数实例化,这里采用正态分布初始化参数
        self.a = torch.nn.Parameter(torch.randn(()))
        self.b = torch.nn.Parameter(torch.randn(()))
        self.c = torch.nn.Parameter(torch.randn(()))
        self.d = torch.nn.Parameter(torch.randn(()))

    def forward(self, x):
        """  定义前向传播（输入x，输出y_pred），这是定义网络的关键,这就是在定义网络的结构 """
        return self.a + self.b + self.c * x ** 2 + self.d * x ** 3

    def string(self):
        """ 这个和普通python class一样可以自定义string方法"""
        return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3'


# 生成数据
x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)
# 实例化模型
model = Polynomial3()
# 构建损失函数和优化器
# 注意这里的reduction='sum'，求损失是对一个batch求损失再求和，反正要求损失最后是一个标量
criterion = torch.nn.MSELoss(reduction='sum')
# 构建优化器，要喂入模型的参数，model.parameters()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)
# 进行训练
for t in range(2000):
    # 前向传播
    y_pred = model(x)
    # 计算loss
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())
    # 清除梯度缓存并反向传播
    optimizer.zero_grad()
    loss.backward()
    # 进行优化
    optimizer.step()
print(f'Result: {model.string()}')

# %%
# 控制流 + 权重共享
# 作为动态图和权重共享的示例，我们实现了一个非常奇怪的模型：
# 这个模型的的阶数是在动态变换的：前3阶0-3被写死了，然后随机加上一个4阶或者5阶或者不加
# 对于此模型，我们可以使用常规的 Python 流控制来实现循环，
# 并且可以通过在定义正向传播时简单地多次重复使用相同的参数来实现权重共享。

import random
import torch
import math

class DynamicNet(torch.nn.Module):
    def __init__(self):
        """
        因为要用到5阶，所以我们生成5个参数
        """
        super(DynamicNet, self).__init__()
        self.a = torch.nn.Parameter(torch.randn(()))
        self.b = torch.nn.Parameter(torch.randn(()))
        self.c = torch.nn.Parameter(torch.randn(()))
        self.d = torch.nn.Parameter(torch.randn(()))
        self.e = torch.nn.Parameter(torch.randn(()))

    def forward(self, x):
        """
        前向传播，我们随机选取4到5层，然后重复使用e来计算
        For the forward pass of the model, we randomly choose either 4, 5
        and reuse the e parameter to compute the contribution of these orders.

        因为每次前向传播都是生成了一个动态计算图，我们可以用python的控制流（循环，条件）来定义模型的前向传播

        Since each forward pass builds a dynamic computation graph, we can use normal
        Python control-flow operators like loops or conditional statements when
        defining the forward pass of the model.



        我们将看到，重复使用相同的参数多次定义一个计算图是相当安全的
        Here we also see that it is perfectly safe to reuse the same parameter many
        times when defining a computational graph.
        """

        # 这里是写死的定义4阶
        y = self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3
        # 这里进行随机，注意python的random.randint 和 numpy torch 的random.randint是不同的
        # numpy是左闭右开半开区间，python是闭区间
        """
        import random as pyton_random
        from numpy import random as numpy_random
        
        for i in range(5):
            # python可以随机出1和2，闭区间
            print('python', pyton_random.randint(1, 2))
            # numpy只能随机出1，左闭右开区间
            print('numpy', numpy_random.randint(1, 2))
        """
        # 这里用的是python的random可以随机出4,5,6再加上range的特性
        # 这个循环可以是：不做循环，exp=4,exp=5三种情况
        for exp in range(4, random.randint(4, 6)):
            y = y + self.e * x ** exp
        return y

    def string(self):
        """
        定义string方法
        """
        return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3 + {self.e.item()} x^4 ? + {self.e.item()} x^5 ?'


# 生成数据
x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)

# 实例化模型
model = DynamicNet()

# 构建我们的损失函数和优化器，训练这个奇怪的模型是困难的，因此我们用带动量的SGD
# Construct our loss function and an Optimizer. Training this strange model with
# vanilla stochastic gradient descent is tough, so we use momentum
criterion = torch.nn.MSELoss(reduction='sum')
# 这里设置了动量
optimizer = torch.optim.SGD(model.parameters(), lr=1e-8, momentum=0.9)
for t in range(30000):
    # Forward pass: Compute predicted y by passing x to the model
    # 前向传播
    y_pred = model(x)

    # Compute and print loss
    # 计算loss并打印
    loss = criterion(y_pred, y)
    if t % 2000 == 1999:
        print(t, loss.item())

    # Zero gradients, perform a backward pass, and update the weights.
    # 清除梯度缓存，并反向传播，并更新参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
# 打印结果
print(f'Result: {model.string()}')
